/**
 * Created by Administrator on 2025/5/28.
 * */
#include "AIPCNet.h"
#include <cuda_runtime.h>
#include <iostream>
#include <thread>

BaseObject *createAIPCNet() {
    return new AIPCNet();
}

void destroyAIPCNet(BaseObject *p) {
    SAFE_DELETE(p);
}

/*__global__ void TaskKernel(char *in, char *out) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    out[tid] = in[tid];
}*/

AIPCNet::AIPCNet(){
    this->setGPUDevice(0);
    prevStep = -1;
    // 订阅数据源主题Topic
    sub = zmq::socket_t(ctx, zmq::socket_type::sub);
    sub.connect(topicAddr.c_str());
    sub.set(zmq::sockopt::subscribe, topicName);
}

AIPCNet::~AIPCNet() {
}

void AIPCNet::onIssueStreamedCudaOperations(buffer_table_t *input, buffer_table_t *output, cudaStream_t stream) {
    std::this_thread::sleep_for(std::chrono::milliseconds(1000));
    int step;
    zmq::message_t msgTopicName;
    zmq::message_t msgData;
    if (!sub.recv(msgTopicName, zmq::recv_flags::none)) {
        std::cerr << "Subscribe: Failed to receive topic name" << std::endl;
    }
    std::string topicName_(static_cast<char*>(msgTopicName.data()), msgTopicName.size());
    if (!sub.recv(msgData, zmq::recv_flags::none)) {
        std::cerr << "Subscribe: Failed to receive data" << std::endl;
    }
    memcpy(&step, msgData.data(), sizeof(float));
    if (prevStep == step) {
        return ;
    }
    int aiPCNetState = GlobalParameters::shareInstance()->get_parameter_aiPCNetState(dynamic_cast<TaskModule*>(this));
    if (aiPCNetState==static_cast<int>(AIPCNetStateEnum::Notify_Initial)) {
        std::cout << "AIPCTrain.AIPCNet: 开始训练模型..." << std::endl;
        auto aiPCNetStateHolder = GlobalParameters::shareInstance()->findParameter("aiPCNetState");
        int *aiPCNetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiPCNetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiPCNetState = static_cast<int>(AIPCNetStateEnum::Initial_Training);
    }
    if (aiPCNetState==static_cast<int>(AIPCNetStateEnum::Initial_Training) && step==STEP3_1) {
        std::cout << "AIPCTrain.AIPCNet: 模型训练完成，通知AINodeInfer::AIPC模块升级模型" << std::endl;
        auto aiPCNetStateHolder = GlobalParameters::shareInstance()->findParameter("aiPCNetState");
        int *aiPCNetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiPCNetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiPCNetState = static_cast<int>(AIPCNetStateEnum::Initial_Trained);
        // 通知AINodeInfer::AIPC模块升级模型
        int aiPCStateVal = 1;
        cudaMemcpyAsync(output->list[1]->data, &aiPCStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
    } else {
        auto aiPCNetStateHolder = GlobalParameters::shareInstance()->findParameter("aiPCNetState");
        int *aiPCNetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiPCNetStateHolder, dynamic_cast<TaskModule*>(this)));
        if (*aiPCNetState == static_cast<int>(AIPCNetStateEnum::Initial_Trained)) {
            // 恢复AIPC
            int aiPCStateVal = 0;
            cudaMemcpyAsync(output->list[1]->data, &aiPCStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
        }
    }
}
